
import com.volcengine.ark.runtime.model.completion.chat.*;
import com.volcengine.ark.runtime.service.ArkService;
import io.reactivex.Flowable;
import io.reactivex.disposables.Disposable;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
// 请确保您已将 API Key 存储在环境变量 ARK_API_KEY 中
// 初始化Ark客户端，从环境变量中读取您的API Key
public class ChatCompletionsVisionExample {
    // 从环境变量中获取您的 API Key。此为默认方式，您可根据需要进行修改
    static String apiKey = "f3447899-2b7c-45cc-afe1-bb9e90ed9b15";
    // 此为默认路径，您可根据业务所在地域进行配置
    static String baseUrl = "https://ark.cn-beijing.volces.com/api/v3";
    static ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
    static Dispatcher dispatcher = new Dispatcher();
    static ArkService service = ArkService.builder().dispatcher(dispatcher).connectionPool(connectionPool).baseUrl(baseUrl).apiKey(apiKey).build();


    public static void main(String[] args) {
        List<ChatMessage> messagesList = new ArrayList<>();

        messagesList.add(
                ChatMessage.builder().role(ChatMessageRole.USER)
                        .content("帮我生成一篇java基础知识的文章，要求1000字左右，并且有代码示例讲解。")
                        .build()
        );

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model("ep-20250308142727-jkxrv")
                .messages(messagesList)
                .build();


        Flowable<ChatCompletionChunk> stream = service.streamChatCompletion(chatCompletionRequest);

        // 订阅流并处理数据
        Disposable disposable = stream
                .doOnError(Throwable::printStackTrace)
                .subscribe(
                        partialResult -> {
                            System.out.print(partialResult.getChoices().get(0).getMessage().getContent());
                        },
                        error -> {
                            System.err.println("发生错误: " + error.getMessage());
                        },
                        () -> {
                            System.out.println("\n流式响应完成");
                        }
                );

        // 等待流处理完成
        try {
            while (!disposable.isDisposed()) {
                Thread.sleep(100);
            }
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        // 关闭服务的执行器
        service.shutdownExecutor();
    }
    /*public static void main(String[] args) {
        System.out.println("----- image input -----");
        final List<ChatMessage> messages = new ArrayList<>();
        final List<ChatCompletionContentPart> multiParts = new ArrayList<>();
        multiParts.add(ChatCompletionContentPart.builder().type("text").text(
                "帮我生成一张图片。"
        ).build());
        *//*multiParts.add(ChatCompletionContentPart.builder().type("image_url").imageUrl(
                new ChatCompletionContentPart.ChatCompletionContentPartImageURL(
                        "https://ark-project.tos-cn-beijing.ivolces.com/images/view.jpeg"
                )
        ).build());*//*

        final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER)
                .multiContent(multiParts).build();
        messages.add(userMessage);

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                // 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
                .model("ep-20250308142727-jkxrv")
                .messages(messages)
                .build();

        service.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> System.out.println(choice.getMessage().getContent()));

        service.shutdownExecutor();
    }*/
}